package cn.b.sky.shop.storage.utils;

import com.google.common.util.concurrent.RateLimiter;

import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;

/**
 * 限制读取速度的InputStream
 * Created by Weijun on 2017/1/11.
 */
public class RateLimitInputStream extends FilterInputStream {

    private RateLimiter rateLimiter;

    public RateLimitInputStream(InputStream in, RateLimiter rateLimiter) {
        super(in);
        this.rateLimiter = rateLimiter;
    }

    /**
     * 限制每秒最多读取多少字节
     *
     * @param in
     * @param bytesPerSecond
     */
    public RateLimitInputStream(InputStream in, long bytesPerSecond) {
        super(in);
        this.rateLimiter = RateLimiter.create(bytesPerSecond);
    }

    @Override
    public int read() throws IOException {
        int ret = super.read();
        if (ret >= 0) {
            rateLimiter.acquire();
        }
        return ret;
    }

    @Override
    public int read(byte[] b, int off, int len) throws IOException {
        int ret = super.read(b, off, len);
        if (ret > 0) {
            rateLimiter.acquire(ret);
        }
        return ret;
    }

    public double getRate() {
        return rateLimiter.getRate();
    }
}
